import torch as pt
import torchvision as ptv
import numpy as np
from thop import profile
from torchsummary import summary
import struct

train_set = ptv.datasets.MNIST("./pytorch_database/mnist/train",train=True,transform=ptv.transforms.ToTensor(),download=True)
test_set = ptv.datasets.MNIST("./pytorch_database/mnist/test",train=False,transform=ptv.transforms.ToTensor(),download=True)

train_dataset = pt.utils.data.DataLoader(train_set,batch_size=100)
test_dataset = pt.utils.data.DataLoader(test_set,batch_size=100)

class MLP(pt.nn.Module):
    def __init__(self):
        super(MLP,self).__init__()
        self.fc1 = pt.nn.Linear(784,512)
        self.fc2 = pt.nn.Linear(512,128)
        self.fc3 = pt.nn.Linear(128,10)
        
        
    def forward(self,din):
        din = din.view(-1,28*28)
        dout = pt.nn.functional.relu(self.fc1(din))
        dout = pt.nn.functional.relu(self.fc2(dout))
        return pt.nn.functional.softmax(self.fc3(dout), dim=1)

model = MLP()
input = pt.randn(1, 28, 28)

# loss func and optim
optimizer = pt.optim.SGD(model.parameters(),lr=0.01,momentum=0.9)
lossfunc = pt.nn.CrossEntropyLoss().cuda()

# accuarcy
def AccuarcyCompute(pred,label):
    pred = pred.cpu().data.numpy()
    label = label.cpu().data.numpy()
#     print(pred.shape(),label.shape())
    test_np = (np.argmax(pred,1) == label)
    test_np = np.float32(test_np)
    return np.mean(test_np)


model = pt.load('./mlp_model.pt')


for x in range(4):
    for i,data in enumerate(train_dataset):
    
        optimizer.zero_grad()
        (inputs,labels) = data
        inputs = pt.autograd.Variable(inputs)
        labels = pt.autograd.Variable(labels)
    
        outputs = model(inputs)
    
        loss = lossfunc(outputs,labels)
        loss.backward()
    
        optimizer.step()
    
        if i % 100 == 0:
            print(i,":",AccuarcyCompute(outputs,labels))


accuarcy_list = []
for i,(inputs,labels) in enumerate(test_dataset):
    inputs = pt.autograd.Variable(inputs).cuda()
    labels = pt.autograd.Variable(labels).cuda()
    outputs = model(inputs)
    accuarcy_list.append(AccuarcyCompute(outputs,labels))
print(sum(accuarcy_list) / len(accuarcy_list))


def save(model, filename):
    
    def traverse(tensor, f):
        for i in tensor:
            if len(i.shape) == 0:
                f.write(struct.pack(">f", (float)(i.data)))
            else:
                traverse(i, f)

    f = open(filename, 'wb')
    for _,param in enumerate(model.named_parameters()):
        traverse(param[1], f)
        print(str(param[0]))

save(model, './model.bin')